#ifndef OMP_EQUITYCALCULATOR_H
#define OMP_EQUITYCALCULATOR_H

#include "CombinedRange.h"
#include "Random.h"
#include "CardRange.h"
#include "HandEvaluator.h"
#include "Constants.h"
#include "Util.h"
#include <chrono>
#include <thread>
#include <mutex>
#include <atomic>
#include <unordered_map>
#include <functional>
#include <array>
#include <cstdint>
#include "WeightedCardRange.h"
#include <mutex>

namespace omp {

    

// Calculates all-in equities in Texas Holdem for given player hand ranges, board cards and dead cards. Supports both
// exact enumeration and monte carlo simulation.
class EquityCalculator
{
public:

    void setEnumerateAll(bool enumerateAll)
    {
        std::lock_guard<std::mutex> lock(mMutex);
        mResults.enumerateAll = enumerateAll;
    }
    EquityCalculator() : mEval() {}

    // 获取每个玩家的牌型胜率（key:牌型字符串，value:胜率）
    std::vector<std::unordered_map<std::string, double>> getTypeEquities() const;
    
    // 计算每个具体手牌组合的胜率
    std::vector<std::unordered_map<std::string, double>> calculateIndividualEquities(const std::vector<WeightedCardRange>& handRanges, uint64_t boardCards = 0, uint64_t deadCards = 0,
                                                                        double stdevTarget = 0.01,
                                                                        double updateInterval = 0.2, unsigned threadCount = 0);
     struct BatchResults {
        std::vector<std::unordered_map<std::string, double>> typeCounts;
        std::vector<std::unordered_map<std::string, double>> typeWins;
        std::vector<uint64_t> wins;
        std::vector<double> ties;
        double winsByPlayerMask[1 << MAX_PLAYERS];
        double winsByPlayerMask_v2[1 << MAX_PLAYERS];
        // 原子数组，用于保存玩家掩码的权重
        // std::array<std::atomic<double>, 1 << MAX_PLAYERS> winsByPlayerMask_v2; 
        BatchResults(unsigned nplayers)
        : typeCounts(nplayers), typeWins(nplayers), wins(nplayers, 0), ties(nplayers, 0.0)
        {
            std::fill(winsByPlayerMask, winsByPlayerMask + (1 << MAX_PLAYERS), 0.0);
            // for (auto& elem : winsByPlayerMask_v2) {
            //         elem.store(0.0, std::memory_order_relaxed);
            //     }
            for (unsigned i = 0; i < nplayers; ++i)
                playerIds[i] = i;
        }
        uint64_t skippedPreflopCombos = 0;
        uint64_t uniquePreflopCombos = 0;
        uint64_t evalCount = 0;
        uint8_t playerIds[MAX_PLAYERS];
    };
    struct Results
    {
        // 增加牌型统计
        std::vector<std::unordered_map<std::string, double>> typeCounts;
        std::vector<std::unordered_map<std::string, double>> typeWins;
        unsigned players = 0;
        double equity[MAX_PLAYERS] = {};
        double weightedTies = 0.0;
        double wins[MAX_PLAYERS] = {};
        double ties[MAX_PLAYERS] = {};
        double winsByPlayerMask[1 << MAX_PLAYERS] = {};
        double winsByPlayerMask_v2[1 << MAX_PLAYERS] = {};  // 新增
        bool isWeightedMonteCarlo = false;  // 新增成员
        Results()
        {
            std::fill(wins, wins + MAX_PLAYERS, 0);
            std::fill(ties, ties + MAX_PLAYERS, 0.0);
            std::fill(winsByPlayerMask, winsByPlayerMask + (1 << MAX_PLAYERS), 0.0);
            std::fill(winsByPlayerMask_v2, winsByPlayerMask_v2 + (1 << MAX_PLAYERS), 0.0);  // 新增
            totalAllSum = 0.0;
            totalWinSqrSum = 0.0;
            hands = intervalHands = 0;
            speed = intervalSpeed = 0;
            time = intervalTime = 0;
            stdev = stdevPerHand = progress = 0;
            preflopCombos = skippedPreflopCombos = evaluatedPreflopCombos = evaluations = 0;
            enumerateAll = finished = isRegularMonteCarlo = false;
        }
        uint64_t hands = 0, intervalHands = 0;
        double speed = 0, intervalSpeed = 0;
        double time = 0, intervalTime = 0;
        double stdev = 0;
        double stdevPerHand = 0;
        double totalAllSum = 0.0;  // 全局累积权重总和
        double totalWinSqrSum = 0.0;  // 累积权重平方和
        double progress = 0;
        uint64_t preflopCombos = 0;
        uint64_t skippedPreflopCombos = 0;
        uint64_t evaluatedPreflopCombos = 0;
        uint64_t evaluations = 0;
        bool enumerateAll = false;
        bool finished = false;
        bool isRegularMonteCarlo = false;
    };

    bool start(const std::vector<CardRange>& handRanges, uint64_t boardCards = 0, uint64_t deadCards = 0,
               bool enumerateAll = false, double stdevTarget = 5e-5,
               std::function<void(const Results&)> callback = nullptr,
               double updateInterval = 0.2, unsigned threadCount = 0);

    bool start_v2(const std::vector<WeightedCardRange>& handRanges, uint64_t boardCards = 0, uint64_t deadCards = 0,
               double stdevTarget = 0.00005,
               std::function<void(const Results&)> callback = nullptr,
               double updateInterval = 0.2, unsigned threadCount = 0);

    void stop()
    {
        mStopped = true;
    }

    void wait()
    {
        for (auto& t : mThreads)
            t.join();
    }

    void setTimeLimit(double seconds)
    {
        std::lock_guard<std::mutex> lock(mMutex);
        mTimeLimit = seconds <= 0 ? INFINITE : seconds;
    }

    void setHandLimit(uint64_t handLimit)
    {
        std::lock_guard<std::mutex> lock(mMutex);
        mHandLimit = handLimit == 0 ? INFINITE : handLimit;
    }

    Results getResults()
    {
        std::lock_guard<std::mutex> lock(mMutex);
        return mUpdateResults;
    }

    const std::vector<CardRange>& handRanges() const
    {
        return mOriginalCardRanges;
    }

private:
    // 在private部分补充声明

    struct HandWithPlayerIdx
    {
        std::array<uint8_t,2> cards;
        unsigned playerIdx;
    };

    void enumerateWeighted();
    void enumerateBoardWeighted(const HandWithPlayerIdx* playerHands, unsigned nplayers,
                            const Hand& board, uint64_t usedCardsMask, double comboWeight, BatchResults* stats);
    void enumerateBoardWeightedRec(const Hand* playerHands, unsigned nplayers, BatchResults* stats,
                                const Hand& board, unsigned* deck, unsigned ndeck, unsigned* suitCounts,
                                unsigned cardsLeft, unsigned start, double weight,
                                const HandWithPlayerIdx* originalPlayerHands);
                        
    typedef XoroShiro128Plus Rng;
    std::string handToType(const std::array<uint8_t, 2>& hand,bool preserveSpecific = false) const;
    // 允许const成员函数修改该成员变量
    mutable std::mutex mMutex;
    static const size_t MAX_LOOKUP_SIZE = 1000000;
    static const size_t MAX_COMBINED_RANGE_SIZE = 10000;
    static const uint64_t INFINITE = ~0ull;

    void simulateRegularMonteCarlo();
    // void simulateWeightedRandomWalkMonteCarlo(const std::vector<WeightedCardRange>&, uint64_t, uint64_t, double);
    void simulateWeightedRegularMonteCarlo(const std::vector<WeightedCardRange>& handRanges, uint64_t boardCards, uint64_t deadCards, double stdevTarget);
    void simulateWeightedRandomWalkMonteCarlo();
    void simulateRandomWalkMonteCarlo();
    bool randomizeHoleCards(uint64_t &usedCardsMask, unsigned* comboIndexes, Hand* playerHands,
                            Rng& rng, FastUniformIntDistribution<unsigned,21>*comboDists);
    OMP_FORCE_INLINE void randomizeBoard(Hand& board, unsigned remainingCards, uint64_t usedCardsMask,
                        Rng& rng, FastUniformIntDistribution<unsigned,16>& cardDist);
    template<bool tFlushPossible = true>
    OMP_FORCE_INLINE void evaluateHands(const Hand* playerHands, unsigned nplayers, const Hand& board,
            BatchResults* stats, unsigned weight);
    OMP_FORCE_INLINE void evaluateHands_v2(const Hand* playerHands, unsigned nplayers, const Hand& board,
            BatchResults* stats, double weight, unsigned* winnersMask);
    void enumerate();
    void enumerateBoard(const HandWithPlayerIdx* playerHands, unsigned nplayers,
                   const Hand& board, uint64_t usedCardsMask, BatchResults* stats);
    void enumerateBoardRec(const Hand* playerHands, unsigned nplayers, BatchResults* stats,
                           const Hand& board, unsigned* deck, unsigned ndeck,  unsigned* suitCounts,
                           unsigned k, unsigned start, unsigned weight, const HandWithPlayerIdx* originalPlayerHands);
    bool lookupResults(uint64_t hash, BatchResults& results);
    bool lookupPrecalculatedResults(uint64_t hash, BatchResults& results) const;
    void storeResults(uint64_t hash, const BatchResults& results);
    static unsigned transformSuits(HandWithPlayerIdx* playerHands, unsigned nplayers,
                                   uint64_t* boardCards, uint64_t* usedCards);
    static uint64_t calculateUniquePreflopId(const HandWithPlayerIdx* playerHands, unsigned nplayers);
    static Hand getBoardFromBitmask(uint64_t board);
    static std::vector<std::vector<std::array<uint8_t,2>>> removeInvalidCombos(const std::vector<CardRange>& handRanges,
                                                               uint64_t reservedCards);
    std::pair<uint64_t,uint64_t> reserveBatch(uint64_t batchCount);
    uint64_t getPreflopCombinationCount();
    uint64_t getPostflopCombinationCount();

    void updateResults(const BatchResults& stats, bool finished);
    double combineResults(const BatchResults& batch);
    double combineResults_v2(const BatchResults& batch);
    void outputLookupTable() const;

    std::vector<std::thread> mThreads;
    std::atomic<bool> mStopped;
    unsigned mUnfinishedThreads;
    std::chrono::high_resolution_clock::time_point mLastUpdate;
    Results mResults, mUpdateResults;
    double mBatchSum, mBatchSumSqr, mBatchCount;
    uint64_t mEnumPosition;
    std::unordered_map<uint64_t, BatchResults> mLookup;
    std::vector<CardRange> mOriginalCardRanges;
    std::vector<WeightedCardRange> mOriginalHandRanges;
    std::vector<std::vector<std::array<uint8_t,2>>> mHandRanges;
    CombinedRange mCombinedRanges[MAX_PLAYERS];
    unsigned mCombinedRangeCount;
    uint64_t mDeadCards, mBoardCards;
    HandEvaluator mEval;
    double mStdevTarget = 5e-5, mTimeLimit = (double)INFINITE, mUpdateInterval = 0.1;
    uint64_t mHandLimit = INFINITE;
    std::function<void(const Results& results)> mCallback;
    static const std::vector<uint64_t> PRECALCULATED_2PLAYER_RESULTS;
};

}

#endif // OMP_EQUITYCALCULATOR_H